
#include "stdafx.h"
#include "ProfilingThread.h"
#include "HookAPI.H"
#include "modulestree.h"
#include "HeapProxy.h"
#include "resource.h"
#include "callstack.h"

HWND	__hStatusWnd = NULL;
HWND	__hConsoleWnd = NULL;
HMODULE __hModDLL = NULL;
DWORD	__dwModSize = 0;
DWORD	__dwHostProcessId = 0;
DWORD	(WINAPI * __HostOutputProc )(LPVOID lpParameter ) = NULL;
DWORD WINAPI __inject_thread( LPVOID lpParameter );
HANDLE __profiler_thread_handle = NULL;
int __signal = 1;
HWND __hwnd_dlg = NULL;

void __update_depends( HMODULE hm,mod_vec_t& mod_vec );
__declspec( thread ) char __print_buf[1024];



void start_profile_thread( HMODULE hmod )
{
#ifdef _DEBUG
	//MessageBox(0,"For Debugger Attaching...","WAIT",0);
#endif
	__profiler_thread_handle = ::CreateThread( NULL,0,__inject_thread,NULL,CREATE_SUSPENDED,NULL );
}

void on_load_library( HMODULE hm )
{
	__update_depends( hm,__mods_vec );
}


HMODULE (WINAPI *IAT_LoadLibraryW)( LPCWSTR lpFileName ) = NULL;
HMODULE WINAPI New_LoadLibraryW( LPCWSTR lpFileName )
{
	HMODULE hm = IAT_LoadLibraryW( lpFileName );

	on_load_library( hm );
	
	return hm;
};

HMODULE (WINAPI *IAT_LoadLibraryA)( LPCSTR lpFileName ) = NULL;

HMODULE WINAPI New_LoadLibraryA( LPCSTR lpFileName )
{
	HMODULE hm = IAT_LoadLibraryA( lpFileName );

	on_load_library( hm );

	return hm;
};

HMODULE (WINAPI *IAT_LoadLibraryExW)( LPCWSTR lpFileName,HANDLE hFile,DWORD dwFlags ) = NULL;
HMODULE WINAPI New_LoadLibraryExW(   
	LPCWSTR lpFileName,
	HANDLE hFile,
	DWORD dwFlags 
	)
{
	HMODULE hm = IAT_LoadLibraryExW( lpFileName,hFile,dwFlags );

	on_load_library( hm );
	
	return hm;
};

HMODULE (WINAPI *IAT_LoadLibraryExA)( 	LPCSTR lpFileName,HANDLE hFile,DWORD dwFlags ) = NULL;

HMODULE WINAPI New_LoadLibraryExA( 
	LPCSTR lpFileName,
	HANDLE hFile,
	DWORD dwFlags
	)
{
	HMODULE hm = IAT_LoadLibraryExA( lpFileName,hFile,dwFlags );

	on_load_library( hm );

	return hm;
};

bool CheckKernelDLL( const char* path,const char*name)
{
	if( stricmp( path,"c:\\windows\\system32" ) == 0)
	{
		if( strstr( name,".dll") )
		{
			return true;
		}
	}

	return false;
}


BOOL CALLBACK __dlg_proc(HWND hwndDlg, UINT message, WPARAM wParam, LPARAM lParam) 
{ 
    BOOL fError; 
	HWND hItem = NULL;

	if( message ==  WM_MPROFILER_MESSAGE )
	{
		int idx = (int) wParam;
		if( idx >=0 )
			set_modules_tree_item( idx,true );
		return TRUE;
	}

 
    switch (message) 
    { 
        case WM_INITDIALOG: 
           // CheckDlgButton(hwndDlg, ID_ABSREL, fRelative); 
			
            return TRUE;
	
		case WM_SIZE:
			hItem = ::GetDlgItem( hwndDlg,IDC_TREE1 );
			RECT rect;
			::GetClientRect( hwndDlg,&rect );
			
			::MoveWindow( hItem,rect.left,rect.top,rect.right - rect.left,rect.bottom - rect.left,TRUE );
			return TRUE;
 
        case WM_COMMAND: 
            switch (LOWORD(wParam)) 
            { 
                case IDOK: 
                    //fRelative = IsDlgButtonChecked(hwndDlg, ID_ABSREL); 
					// iLine = GetDlgItemInt(hwndDlg, ID_LINE, &fError, fRelative); 
					// if (fError) 
					//{ 
                    //    MessageBox(hwndDlg, SZINVALIDNUMBER, SZGOTOERR, MB_OK); 
                    //    SendDlgItemMessage(hwndDlg, ID_LINE, EM_SETSEL, 0, -1L); 
                    //} 
                    //else 

                    // Notify the owner window to carry out the task. 
 
                    return TRUE; 
 
                case IDCANCEL: 
                    DestroyWindow(hwndDlg); 
                   // hwndGoto = NULL; 
                    return TRUE; 
            } 
    } 
    return FALSE; 
} 

void update_hooks( HMODULE hm,mod_vec_t& mod_vec )
{
	HMODULE hKernel = ::GetModuleHandleA("kernel32.dll");

	for( int i = 0;i<__mods_vec.size();i++ )
	{
		if( !__mods_vec[i].is_reference )
		{
			if( CheckKernelDLL( __mods_vec[i].path,__mods_vec[i].name) )
				continue;
		
			if( __mods_vec[i].instance == __hModDLL )
				continue;
		
			//BOOL br = HookImportedFunc( __mods_vec[i].instance,"kernel32.dll","LoadLibraryW",(PROC)New_LoadLibraryW,(PROC*)&IAT_LoadLibraryW );
			//if( !br )
			{
				HOOKAPI api = {"LoadLibraryW",(DWORD)GetProcAddress(hKernel,"LoadLibraryW"),(PROC)New_LoadLibraryW,0};
				HookIAT( __mods_vec[i].instance,"Kernel32.dll",&api );
				if( !IAT_LoadLibraryW )
					IAT_LoadLibraryW = (HMODULE (__stdcall *)(LPCWSTR))api.old_proc;
			}

			{
				HOOKAPI api = {"LoadLibraryA",(DWORD)GetProcAddress(hKernel,"LoadLibraryA"),(PROC)New_LoadLibraryA,0};

				HookIAT( __mods_vec[i].instance,"Kernel32.dll",&api );

				if( !IAT_LoadLibraryA )
					IAT_LoadLibraryA = (HMODULE (__stdcall *)(LPCSTR))api.old_proc;
			}

			{
				HOOKAPI api = {"LoadLibraryExW",(DWORD)GetProcAddress(hKernel,"LoadLibraryExW"),(PROC)New_LoadLibraryExW,0};
				HookIAT( __mods_vec[i].instance,"Kernel32.dll",&api );
				if( !IAT_LoadLibraryExW )
					IAT_LoadLibraryExW = (HMODULE (__stdcall *)(LPCWSTR,HANDLE,DWORD))api.old_proc;
			}

			{
				HOOKAPI api = {"LoadLibraryExA",(DWORD)GetProcAddress(hKernel,"LoadLibraryExA"),(PROC)New_LoadLibraryExA,0};

				HookIAT( __mods_vec[i].instance,"Kernel32.dll",&api );

				if( !IAT_LoadLibraryExA )
					IAT_LoadLibraryExA = (HMODULE (__stdcall *)(LPCSTR,HANDLE,DWORD))api.old_proc;
			}
		}
	}

	for( int i = 0;i<__mods_vec.size();i++ )
	{
		if( !__mods_vec[i].is_reference )
		{
			if( CheckKernelDLL( __mods_vec[i].path, __mods_vec[i].name ) )
				continue;

			if( __mods_vec[i].instance == __hModDLL )
				continue;

			HookHeapAPIs( __mods_vec[i].instance )	;
		}
	}
}

void __update_depends( HMODULE hm,mod_vec_t& mod_vec )
{
	DWORD dwProcessId = ::GetCurrentProcessId();
	HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                            PROCESS_VM_READ,
                            FALSE, dwProcessId );
	if( hm )
	{
		UpdateModule( hm,mod_vec );
	}
	else
	{
		GetImportModules( NULL,NULL,0,mod_vec,-1 );
	}
	
	CreateModules( hProcess,mod_vec,false );
	clear_modules_tree();
	create_modules_tree( mod_vec );
	expand_modules_tree();
	update_hooks( hm,mod_vec );

	CloseHandle( hProcess );
}

void resume_profile_thread ()
{
	

	while( __signal )
	{
		::ResumeThread( __profiler_thread_handle );
		Sleep(50);
	}
}

UINT WM_MPROFILER_MESSAGE = RegisterWindowMessage("mprofiler");

DWORD WINAPI __inject_thread( LPVOID lpParameter )
{
	init_callstack();

	LoadLibrary("RICHED20.DLL");

	INITCOMMONCONTROLSEX InitCtrls;
	InitCtrls.dwSize = sizeof(InitCtrls);
	InitCtrls.dwICC = ICC_WIN95_CLASSES;
	InitCommonControlsEx(&InitCtrls);

	__hwnd_dlg = ::CreateDialogA( __hModDLL,MAKEINTRESOURCE(IDD_DIALOG1),NULL,__dlg_proc );

	HWND hwnd_tree_ctrl = GetDlgItem(__hwnd_dlg,IDC_TREE1);
	init_tree_ctrl( hwnd_tree_ctrl );
	
	char name[MAX_PATH];
	::GetModuleFileNameA( NULL, name,MAX_PATH );

	::SetWindowTextA( __hwnd_dlg,name );

	ShowWindow( __hwnd_dlg,SW_SHOW );
	
	HWND hConsole = GetConsoleWindow();
	if( !hConsole )
	{
		AllocConsole();
		freopen( "CON","w",stdout );
	}

	HMODULE hKernel32 = ::GetModuleHandleA( "kernel32.dll" );
	PROC pLoadLibraryW = ::GetProcAddress( hKernel32,"LoadLibraryW" );
	HMODULE hMod = NULL;
	::GetModuleHandleExA( GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,(LPCSTR)pLoadLibraryW,&hMod );
	
	if( hMod == hKernel32 )
	{
		__update_depends( NULL,__mods_vec );
	}	

	__signal = 0;

	BOOL bRet;
	MSG msg;

	if( __hwnd_dlg )
	{
		while ((bRet = GetMessage(&msg, NULL, 0, 0)) != 0) 
		{ 	
			TranslateMessage(&msg); 
			DispatchMessage(&msg); 	
		} 
	}

	clean_callstack();

	return 0;
}